{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Load Nanodet in PyTorch\n",
    "\n",
    "This article is an introductory tutorial to infer nanodet with PyTorch.\n",
    "\n",
    "**Note**, we suppose this notebook is in the root directory of nonadet!\n",
    "\n",
    "## Install nanodet first\n",
    "\n",
    "For us to begin with, PyTorch should be installed. TorchVision is also required since we will be using it as our model zoo.\n",
    "\n",
    "A quick solution is to install via pip\n",
    "\n",
    "```shell\n",
    "pip install torch==1.7.1\n",
    "pip install torchvision==0.8.2\n",
    "```\n",
    "\n",
    "or please refer to official site https://pytorch.org/get-started/locally/\n",
    "\n",
    "PyTorch versions should be backwards compatible but should be used with the proper TorchVision version.\n",
    "\n",
    "And then don't forget to install other dependencies.\n",
    "\n",
    "```shell\n",
    "pip install -r requirements.txt\n",
    "```\n",
    "\n",
    "Next enter the key part, let's install `nanodet`!\n",
    "\n",
    "```shell\n",
    "python setup.py develop\n",
    "```\n",
    "\n",
    "## Set Environmental Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import torch\n",
    "\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "\n",
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.backends.cudnn.enabled = True\n",
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Model Configuration and Logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nanodet.util import cfg, load_config, Logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = 'config/nanodet-m.yml'\n",
    "model_path = 'workspace/nanodet_m/nanodet_m.pth'\n",
    "image_path = 'demo_mnn/imgs/000252.jpg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'config/nanodet-m.yml'",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mFileNotFoundError\u001B[0m                         Traceback (most recent call last)",
      "\u001B[1;32m<ipython-input-11-7208a70fa256>\u001B[0m in \u001B[0;36m<module>\u001B[1;34m\u001B[0m\n\u001B[1;32m----> 1\u001B[1;33m \u001B[0mload_config\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mcfg\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mconfig_path\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0m\u001B[0;32m      2\u001B[0m \u001B[0mlogger\u001B[0m \u001B[1;33m=\u001B[0m \u001B[0mLogger\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;33m-\u001B[0m\u001B[1;36m1\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0muse_tensorboard\u001B[0m\u001B[1;33m=\u001B[0m\u001B[1;32mFalse\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m      3\u001B[0m \u001B[1;33m\u001B[0m\u001B[0m\n",
      "\u001B[1;32m~\\Desktop\\nanodet-main\\nanodet-main\\nanodet\\util\\config.py\u001B[0m in \u001B[0;36mload_config\u001B[1;34m(cfg, args_cfg)\u001B[0m\n\u001B[0;32m     29\u001B[0m \u001B[1;32mdef\u001B[0m \u001B[0mload_config\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mcfg\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0margs_cfg\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m:\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m     30\u001B[0m     \u001B[0mcfg\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mdefrost\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[1;32m---> 31\u001B[1;33m     \u001B[0mcfg\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mmerge_from_file\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0margs_cfg\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0m\u001B[0;32m     32\u001B[0m     \u001B[0mcfg\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mfreeze\u001B[0m\u001B[1;33m(\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m     33\u001B[0m \u001B[1;33m\u001B[0m\u001B[0m\n",
      "\u001B[1;32m~\\Desktop\\nanodet-main\\nanodet-main\\nanodet\\util\\yacs.py\u001B[0m in \u001B[0;36mmerge_from_file\u001B[1;34m(self, cfg_filename)\u001B[0m\n\u001B[0;32m    209\u001B[0m     \u001B[1;32mdef\u001B[0m \u001B[0mmerge_from_file\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mself\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mcfg_filename\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m:\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m    210\u001B[0m         \u001B[1;34m\"\"\"Load a yaml config file and merge it this CfgNode.\"\"\"\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[1;32m--> 211\u001B[1;33m         \u001B[1;32mwith\u001B[0m \u001B[0mopen\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mcfg_filename\u001B[0m\u001B[1;33m,\u001B[0m \u001B[1;34m\"r\"\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mencoding\u001B[0m\u001B[1;33m=\u001B[0m\u001B[1;34m'utf-8'\u001B[0m\u001B[1;33m)\u001B[0m \u001B[1;32mas\u001B[0m \u001B[0mf\u001B[0m\u001B[1;33m:\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0m\u001B[0;32m    212\u001B[0m             \u001B[0mcfg\u001B[0m \u001B[1;33m=\u001B[0m \u001B[0mself\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mload_cfg\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mf\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0;32m    213\u001B[0m         \u001B[0mself\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mmerge_from_other_cfg\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mcfg\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n",
      "\u001B[1;31mFileNotFoundError\u001B[0m: [Errno 2] No such file or directory: 'config/nanodet-m.yml'"
     ]
    }
   ],
   "source": [
    "load_config(cfg, config_path)\n",
    "logger = Logger(-1, use_tensorboard=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Model Predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from demo.demo import Predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model size is  1.0x\n",
      "init weights...\n",
      "=> loading pretrained model https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth\n",
      "Finish initialize Lite GFL Head.\n"
     ]
    }
   ],
   "source": [
    "predictor = Predictor(cfg, model_path, logger, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "forward time: 0.061s | decode time: 0.009s | "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:3063: UserWarning: Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0. Please specify align_corners=True if the old behavior is desired. See the documentation of nn.Upsample for details.\n",
      "  \"See the documentation of nn.Upsample for details.\".format(mode))\n"
     ]
    }
   ],
   "source": [
    "meta, res = predictor.inference(image_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Detection output visualisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nanodet.util import overlay_bbox_cv\n",
    "\n",
    "from IPython.display import display\n",
    "from PIL import Image\n",
    "\n",
    "def cv2_imshow(a, convert_bgr_to_rgb=True):\n",
    "    \"\"\"A replacement for cv2.imshow() for use in Jupyter notebooks.\n",
    "    Args:\n",
    "        a: np.ndarray. shape (N, M) or (N, M, 1) is an NxM grayscale image. shape\n",
    "            (N, M, 3) is an NxM BGR color image. shape (N, M, 4) is an NxM BGRA color\n",
    "            image.\n",
    "        convert_bgr_to_rgb: switch to convert BGR to RGB channel.\n",
    "    \"\"\"\n",
    "    a = a.clip(0, 255).astype('uint8')\n",
    "    # cv2 stores colors as BGR; convert to RGB\n",
    "    if convert_bgr_to_rgb and a.ndim == 3:\n",
    "        if a.shape[2] == 4:\n",
    "            a = cv2.cvtColor(a, cv2.COLOR_BGRA2RGBA)\n",
    "        else:\n",
    "            a = cv2.cvtColor(a, cv2.COLOR_BGR2RGB)\n",
    "    display(Image.fromarray(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'overlay_bbox_cv' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "\u001B[1;32m<ipython-input-4-b3215422e92e>\u001B[0m in \u001B[0;36m<module>\u001B[1;34m\u001B[0m\n\u001B[1;32m----> 1\u001B[1;33m \u001B[0mresult\u001B[0m \u001B[1;33m=\u001B[0m \u001B[0moverlay_bbox_cv\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mmeta\u001B[0m\u001B[1;33m[\u001B[0m\u001B[1;34m'raw_img'\u001B[0m\u001B[1;33m]\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mres\u001B[0m\u001B[1;33m[\u001B[0m\u001B[1;36m0\u001B[0m\u001B[1;33m]\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mcfg\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mclass_names\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mscore_thresh\u001B[0m\u001B[1;33m=\u001B[0m\u001B[1;36m0.35\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0m\u001B[0;32m      2\u001B[0m \u001B[1;33m\u001B[0m\u001B[0m\n",
      "\u001B[1;31mNameError\u001B[0m: name 'overlay_bbox_cv' is not defined"
     ]
    }
   ],
   "source": [
    "result = overlay_bbox_cv(meta['raw_img'], res[0], cfg.class_names, score_thresh=0.35)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'cv2_imshow' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "\u001B[1;32m<ipython-input-3-01591af27675>\u001B[0m in \u001B[0;36m<module>\u001B[1;34m\u001B[0m\n\u001B[0;32m      1\u001B[0m \u001B[0mimshow_scale\u001B[0m \u001B[1;33m=\u001B[0m \u001B[1;36m1.0\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[1;32m----> 2\u001B[1;33m \u001B[0mcv2_imshow\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mcv2\u001B[0m\u001B[1;33m.\u001B[0m\u001B[0mresize\u001B[0m\u001B[1;33m(\u001B[0m\u001B[0mresult\u001B[0m\u001B[1;33m,\u001B[0m \u001B[1;32mNone\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mfx\u001B[0m\u001B[1;33m=\u001B[0m\u001B[0mimshow_scale\u001B[0m\u001B[1;33m,\u001B[0m \u001B[0mfy\u001B[0m\u001B[1;33m=\u001B[0m\u001B[0mimshow_scale\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m)\u001B[0m\u001B[1;33m\u001B[0m\u001B[1;33m\u001B[0m\u001B[0m\n\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m: name 'cv2_imshow' is not defined"
     ]
    }
   ],
   "source": [
    "imshow_scale = 1.0\n",
    "cv2_imshow(cv2.resize(result, None, fx=imshow_scale, fy=imshow_scale))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "YOLOv5 Tutorial",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}